package com.seanLab.tool.TagSuggestModel;

import com.seanLab.domain.TagInfo;
import com.seanLab.dto.SuggestModelArticleDto;
import com.seanLab.tool.TagSuggestModel.ExpandRank.ExpandRankKE;
import com.seanLab.tool.TagSuggestModel.ExpandRank.TrainExpandRank;
import com.seanLab.tool.TagSuggestModel.ExpandRank.common.ConfigIO;
import org.thunlp.misc.WeightString;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class FeatureExtractorByExpandRank extends FeatureExtractor {
    private static final String WORK_PATH = "workingDir";
    private static final String CONFIG = "num_tags=10;norm=all_log;k=5;minwordfreq=10;mintagfreq=10";
    ExpandRankKE expandRankKE = null;

    public FeatureExtractorByExpandRank() {

    }

    public FeatureInTags extractFeature(RawInfo rawInfo, RawInfoType type) {
        if (expandRankKE == null) {
            try {
                throw new Exception("ExpandRank Model haven't bean loaded.");
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        switch (type) {
            case TITLE:
            case CONTENT:
            case PARA_ABOVE_IMAGE:
            case PARA_UNDER_IMAGE:
            case DESC_OF_IMAGE:
                SuggestModelArticleDto post = new SuggestModelArticleDto();
                post.setContent(((RawInfoInText)rawInfo).getContent());
                if (post.getContent() == null || post.getContent().equals("")) {
                    return new FeatureInTags();
                }
                List<WeightString> tagsTMP = expandRankKE.suggest(post);
                List<TagInfo> tags = new ArrayList<>();
                for (WeightString weightString : tagsTMP) {
                    tags.add(new TagInfo(null, weightString.getText(), null, weightString.getWeight(), 0));
                }
                FeatureInTags ret = new FeatureInTags();
                ret.setFeature(tags);
                return ret;
            case SEMANTIC_OF_IMAGE:
            case TAG:
            default:
                return new FeatureInTags();
        }
    }

    public String trainExpandRank(List<SuggestModelArticleDto> articles) {
        String modelPath;
        TrainExpandRank trainer = new TrainExpandRank();
        try {
            modelPath =  trainer.train(articles,
                    System.getProperty("working_dir", WORK_PATH),
                    ConfigIO.configFromString(System.getProperty("expand_rank_config", CONFIG)));
            loadExpandRankModel(modelPath);
            return modelPath;
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    public void loadExpandRankModel(String modelPath) throws IOException {
        expandRankKE = new ExpandRankKE();
        expandRankKE.loadModel(modelPath);
        expandRankKE.setConfig(ConfigIO.configFromString(System.getProperty("expand_rank_config", CONFIG)));
    }

}
